# LangGraph Architecture Documentation

## Overview

This project implements a **Supervisor Pattern** using LangGraph's `StateGraph` to create an AI assistant with RAG capabilities, tool execution, and multi-layer safety guardrails.

```
                                    +-----------------+
                                    |     START       |
                                    +-----------------+
                                            |
                                            v
                                +------------------------+
                                |   input_guardrails    |
                                |   (Security Check)     |
                                +------------------------+
                                    |           |
                            PASSED  |           | BLOCKED
                                    v           v
                            +-----------+  +-----------+
                            | supervisor|  |  blocked  |---> END
                            |  (Brain)  |  +-----------+
                            +-----------+
                                    |
                    +---------------+---------------+---------------+
                    |               |               |               |
                    v               v               v               v
            +-----------+   +-----------+   +-----------+   +-----------+
            | calculate |   | summarize |   |  search   |   | guardrails|
            +-----------+   +-----------+   +-----------+   +-----------+
                    |               |               |               |
                    +-------+-------+-------+               +-------+-------+
                            |                               |               |
                            v                           PASSED          BLOCKED
                    +-----------+                           |               |
                    | supervisor|<--------------------------+               v
                    +-----------+                                   +-----------+
                                                                    |  blocked  |
                                                                    +-----------+
                                                                           |
                            +-----------+                                  v
                            |  deliver  |--------------------------------> END
                            +-----------+
```

---

## 1. State Definition

### AgentState (TypedDict)

The state that flows through the graph is defined using `TypedDict` with `Annotated` types for proper accumulation.

```python
class AgentState(TypedDict):
    # Accumulates chat history (uses operator.add reducer)
    messages: Annotated[list[BaseMessage], operator.add]

    # Current processed query
    current_query: str

    # Classified intent (calculate, search, summarize, respond)
    intent: str

    # Output from tools
    tool_result: str

    # Raw LLM response before safety check
    response_draft: str

    # Safety check status
    guardrail_passed: bool

    # Reason for blocking (if any)
    guardrail_reason: str

    # Trace of decisions/steps taken (uses operator.add reducer)
    steps: Annotated[list[str], operator.add]
```

### Key Concepts

| Field | Type | Reducer | Purpose |
|-------|------|---------|---------|
| `messages` | `list[BaseMessage]` | `operator.add` | Accumulates conversation history |
| `steps` | `list[str]` | `operator.add` | Accumulates execution trace |
| `guardrail_passed` | `bool` | None | **Must be explicitly set** - defaults to `False` for security |

> **Important**: Fields with `operator.add` reducer will accumulate values across turns. Fields without reducers will be overwritten.

---

## 2. Nodes

### 2.1 input_guardrails

**Purpose**: First security checkpoint - validates user input before processing.

```python
def node_input_guardrails(state: AgentState) -> Dict:
    """
    Checks:
    - Toxicity (banned words)
    - Prompt injection patterns
    - PII detection and redaction
    """
```

**Returns**:
- `guardrail_passed: True/False`
- `guardrail_reason: str` (if blocked)
- `current_query: str` (sanitized version)
- `steps: [...]`

---

### 2.2 supervisor

**Purpose**: The "brain" - decides whether to call a tool or respond directly.

```python
def node_supervisor(state: AgentState, llm: ChatOllama, system_prompt: str) -> Dict:
    """
    Decision Logic:
    1. If needs information -> CHOOSE A TOOL
    2. If has tool result -> USE IT to answer
    3. If simple greeting/chat -> ANSWER directly
    """
```

**Tool Detection**:
```python
TOOL_MAPPING = {
    "calculate": "calculate",
    "calculator": "calculate",
    "calc": "calculate",
    "math": "calculate",
    "search": "search",
    "rag": "search",
    "find": "search",
    "lookup": "search",
    "summarize": "summarize",
    "summary": "summarize",
}
```

**Returns**:
- `intent: str` (calculate, summarize, search, respond)
- `response_draft: str` (if responding)
- `current_query: str`
- `steps: [...]`

---

### 2.3 Tool Nodes

#### calculate
```python
def node_run_calculator(state: AgentState) -> Dict:
    """Executes math calculations using calculate_square tool."""
```

#### summarize
```python
def node_run_summary(state: AgentState) -> Dict:
    """Executes text summarization using generate_text_summary tool."""
```

#### search
```python
def node_run_search(state: AgentState) -> Dict:
    """Executes RAG search using FAISS vector store."""
```

**All tools return**:
- `tool_result: str`
- `steps: [...]`

---

### 2.4 guardrails

**Purpose**: Output security checkpoint - validates LLM response before delivery.

```python
def node_guardrails(state: AgentState) -> Dict:
    """
    Checks:
    - Banned words in output
    - PII leakage
    - Content policy violations
    """
```

**Returns**:
- `guardrail_passed: True/False`
- `guardrail_reason: str` (if blocked)
- `response_draft: str` (sanitized)
- `steps: [...]`

---

### 2.5 deliver

**Purpose**: Final delivery node - creates the AIMessage for the user.

```python
def node_deliver(state: AgentState) -> Dict:
    """Returns final message to user."""
    return {"messages": [AIMessage(content=state["response_draft"])]}
```

---

### 2.6 blocked

**Purpose**: Safety fallback - returns a blocked message when guardrails fail.

```python
def node_blocked(state: AgentState) -> Dict:
    """Returns blocked message with reason."""
    return {"messages": [AIMessage(content=f"Response Blocked: {state['guardrail_reason']}")]}
```

---

## 3. Routing Functions

### 3.1 route_input_guardrails

```python
def route_input_guardrails(state: AgentState) -> Literal["supervisor", "blocked"]:
    """
    Routes after input validation.
    IMPORTANT: Default is False (block) for security.
    """
    if state.get("guardrail_passed", False):  # Default False!
        return "supervisor"
    return "blocked"
```

### 3.2 route_intent

```python
def route_intent(state: AgentState) -> Literal["calculate", "summarize", "search", "respond"]:
    """Routes based on supervisor's decision."""
    intent = state.get("intent", "respond")
    if intent == "calculate": return "calculate"
    if intent == "summarize": return "summarize"
    if intent == "search": return "search"
    return "respond"
```

### 3.3 route_guardrails

```python
def route_guardrails(state: AgentState) -> Literal["deliver", "blocked"]:
    """
    Routes after output validation.
    IMPORTANT: Default is False (block) for security.
    """
    if state.get("guardrail_passed", False):  # Default False!
        return "deliver"
    return "blocked"
```

---

## 4. Graph Construction

```python
def create_agent(personality, memory, model_name, config):
    # 1. Initialize LLM
    llm = create_llm(model_name, config=config)
    sys_prompt = build_system_prompt(personality, config)

    # 2. Create wrapper for supervisor (captures llm and prompt)
    def _supervisor(s):
        return node_supervisor(s, llm, sys_prompt)

    # 3. Build Graph
    graph = StateGraph(AgentState)

    # 4. Add Nodes
    graph.add_node("input_guardrails", node_input_guardrails)
    graph.add_node("supervisor", _supervisor)
    graph.add_node("calculate", node_run_calculator)
    graph.add_node("summarize", node_run_summary)
    graph.add_node("search", node_run_search)
    graph.add_node("guardrails", node_guardrails)
    graph.add_node("deliver", node_deliver)
    graph.add_node("blocked", node_blocked)

    # 5. Set Entry Point
    graph.set_entry_point("input_guardrails")

    # 6. Add Conditional Edges
    graph.add_conditional_edges(
        "input_guardrails",
        route_input_guardrails,
        {"supervisor": "supervisor", "blocked": "blocked"}
    )

    graph.add_conditional_edges(
        "supervisor",
        route_intent,
        {
            "calculate": "calculate",
            "summarize": "summarize",
            "search": "search",
            "respond": "guardrails"
        }
    )

    graph.add_conditional_edges(
        "guardrails",
        route_guardrails,
        {"deliver": "deliver", "blocked": "blocked"}
    )

    # 7. Add Tool -> Supervisor Loops (CYCLIC)
    graph.add_edge("calculate", "supervisor")
    graph.add_edge("summarize", "supervisor")
    graph.add_edge("search", "supervisor")

    # 8. Add End Edges
    graph.add_edge("deliver", END)
    graph.add_edge("blocked", END)

    # 9. Compile with Checkpointer and Recursion Limit
    return graph.compile(
        checkpointer=memory
    ).with_config({"recursion_limit": 25})
```

---

## 5. Checkpointing & Memory

### MemorySaver

The project uses LangGraph's `MemorySaver` for in-memory state persistence.

```python
from langgraph.checkpoint.memory import MemorySaver

# Initialize once per session
if "memory" not in st.session_state:
    st.session_state.memory = MemorySaver()

# Pass to agent
agent = create_agent(
    personality=current_personality,
    memory=st.session_state.memory,
    ...
)
```

### Thread Isolation

Each conversation uses a unique `thread_id` for isolation:

```python
# Generate unique thread
thread_id = str(uuid.uuid4())

# Invoke with thread config
response = agent.invoke(
    {"messages": [HumanMessage(content=query)]},
    config={"configurable": {"thread_id": thread_id}}
)

# Get state for a specific thread
state = agent.get_state({"configurable": {"thread_id": thread_id}})
```

---

## 6. Execution Flow

### Typical Flow (with tool)

```
1. User: "What is the square of 5?"
2. input_guardrails: PASSED
3. supervisor: Decides TOOL: CALCULATE
4. calculate: Returns "25"
5. supervisor: Decides to ANSWER using tool result
6. guardrails: PASSED
7. deliver: "The square of 5 is 25"
8. END
```

### Flow with Blocked Input

```
1. User: "[malicious prompt injection]"
2. input_guardrails: BLOCKED (injection detected)
3. blocked: "Response Blocked: Potential prompt injection detected"
4. END
```

---

## 7. Delta Steps Tracking

To track only the steps from the current turn (not accumulated):

```python
def chat(agent, query: str, thread_id: str):
    # 1. Get state BEFORE invocation
    initial_state = agent.get_state({"configurable": {"thread_id": thread_id}})
    start_count = len(initial_state.values.get("steps", [])) if initial_state.values else 0

    # 2. Invoke agent
    response = agent.invoke(...)

    # 3. Calculate delta
    final_steps = response.get("steps", [])
    delta_steps = final_steps[start_count:]  # Only new steps!

    return result, response, delta_steps
```

---

## 8. Security Considerations

### Default Values

All routing functions use **secure defaults**:

```python
# WRONG (insecure):
if state.get("guardrail_passed", True):  # Assumes safe!

# CORRECT (secure):
if state.get("guardrail_passed", False):  # Assumes unsafe!
```

### Recursion Limit

Prevents infinite loops:

```python
graph.compile(...).with_config({"recursion_limit": 25})
```

### Error Handling

Supervisor has try/except to prevent crashes:

```python
def node_supervisor(...):
    try:
        # ... logic
    except Exception as e:
        logger.error(f"Supervisor error: {e}")
        return {
            "intent": "respond",
            "response_draft": "Error processing request.",
            ...
        }
```

---

## 9. Best Practices

### DO

1. **Always set `guardrail_passed` explicitly** in guardrail nodes
2. **Use `operator.add` reducer** for accumulating fields (messages, steps)
3. **Capture state BEFORE invoke** for delta calculation
4. **Use tool mapping** instead of string matching
5. **Set recursion_limit** to prevent infinite loops
6. **Wrap supervisor in try/except**

### DON'T

1. **Don't assume default True** for security fields
2. **Don't create new MemorySaver** when switching conversations
3. **Don't parse tools with simple string matching**
4. **Don't forget to log errors** in node functions

---

## 10. Diagram Generation

### Static Graph

```python
from src.utils.visualizer import generate_graph_image

# Generate full architecture diagram
generate_graph_image(agent, "data/images/generated_graph.png")
```

### Execution Trace

```python
from src.utils.visualizer import generate_execution_trace_image

# Generate trace for specific execution
generate_execution_trace_image(delta_steps, "data/images/execution_trace_last.png")
```

---

## 11. Troubleshooting

### RAG Search Not Working

1. **Check if index is empty**:
   - Look for logs: `📊 Index contains X vectors`
   - If 0 vectors, re-index documents in the RAG tab

2. **LLM not choosing SEARCH tool**:
   - Check supervisor logs: `🤖 RAW LLM OUTPUT`
   - If no `TOOL: SEARCH`, the prompt may need adjustment
   - Verify the query is about company/documents

3. **Threshold too high**:
   - Check `config/model_config.yaml` → `rag.score_threshold`
   - Default is 0.05, try lowering to 0.01

4. **Debug Steps**:
   ```
   🛡️ [INPUT GUARDRAILS] → Should show ✅ PASSED
   🧠 [SUPERVISOR] → Should show TOOL: SEARCH for info queries
   🔎 [TOOL NODE] → Should show Found X documents
   ```

### Common Issues

| Symptom | Cause | Fix |
|---------|-------|-----|
| "No response generated" | LLM returned empty | Check Ollama is running |
| Always goes to "blocked" | `guardrail_passed` not set | Check input_guardrails logs |
| RAG returns empty | Threshold too high | Lower `score_threshold` |
| Infinite loop | Supervisor always picks tool | Check `recursion_limit` |

---

## References

- [LangGraph Documentation](https://langchain-ai.github.io/langgraph/)
- [LangGraph Concepts](https://langchain-ai.github.io/langgraph/concepts/)
- [StateGraph API](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.StateGraph)
- [Checkpointing](https://langchain-ai.github.io/langgraph/concepts/persistence/)
